import re
import json

from typing import Dict, Any
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.prompts import MessagesPlaceholder
from langchain_core.output_parsers import StrOutputParser


class ExtractManager:
    def __init__(self, cfg: Dict[str, Any]):
        self.retry = cfg.get('retry')
        
        model = ChatOpenAI(
            model=cfg.get('extract_model'),
            base_url=cfg.get('base_url'),
            api_key=cfg.get('api_key'),
            temperature=0.0
        )

        prompt = ChatPromptTemplate.from_messages([
            (
                "system",
                "你是一个信息抽取助手。"
                "请从用户的原始输入中提取以下信息，并严格只输出合法的 JSON：\n\n"
                "1. static：静态信息，用中文 key–value 对表示（如生日、身份证号等，不会变化的属性）；\n"
                "2. dynamic：动态信息，用中文 key–value 对表示（如文件路径、当前任务等会变化的状态）；\n"
                "3. events：一个事件列表，每个事件都要包含：\n"
                "     - time(str): 时间；\n"
                "     - location(str): 地点；\n"
                "     - people(str): 相关人物；\n"
                "     - description(str) 事件描述。\n"
                "4. solutions：问题-解决方案键值对列表，结构为：\n"
                "     - problem(str)：用户遇到的问题或需求；\n"
                "     - solution(str)：对应的解决方案；\n"
                "     - iswork(bool)：用户是否明确反馈解决方案是否有效，有效为true。\n"
                "   当用户最近一次发言是提问，iswork为false；当用户最近一次发言反馈刚才的方法有效，iswork为true。 \n"
                "5. ideas：思想观点（用户画像），以主题(theme)为key，观点(opinion)为value。\n\n"
                "注意：key和value要准确描述信息，可以增加定语限制，如果是属性最好加上是谁的属性，比如“用户的xx”,也可以不使用用户说的原话；\n"
                "      时间可以用“今天”等模糊词；如events缺某key，置空；你要称对话人为“用户”。\n\n"
                "输出JSON示例：\n"
                "{{\n"
                "  \"static\": {{\"用户的生日\": \"2025-05-09\"}},\n"
                "  \"dynamic\": {{\"用户的AA项目的xx文件路径\": \"C:/path/to/file\"}},\n"
                "  \"events\": [\n"
                "    {{\"time\": \"今天\", \"location\": \"商场\", \"people\": \"用户\", \"description\": \"去商场买裤子\"}}\n"
                "  ],\n"
                "  \"solutions\": [\n"
                "    {{\"problem\": \"在使用xx库的时候报xx错误\", \"solution\": \"正确写法是xxx\", \"iswork\": \"true\"}},\n"
                "    {{\"problem\": \"pip xx的时候失败\", \"solution\": \"使用vpn或者清华源\", \"iswork\": \"false\"}}\n"
                "  ],\n"
                "  \"ideas\": {{\n"
                "      \"对于资本剥削劳动者的态度\": \"用户觉得这帮资本家真该死啊\", \n"
                "      \"对党\": \"忠诚！\"\n"
                "  }}\n"
                "}}\n"
                "如果某一类别没有任何信息，请不要在输出的JSON里出现该类别，只输出实际抽取得到的类别字段。\n"
                "注意！一定要提取准确的信息，不确定的不要写！一定不要写一些虚假的信息！\n"
                "注意！不要回复用户的问题，仅做信息提取！\n"
                "以下是最近的对话历史（供你理解背景）：\n"
                "{history}"
            ),
            ("user", "{user_input}")
        ])

        parser = StrOutputParser()

        self.chain = prompt | model | parser

    def extract(self, messages: str):
        n = 2
        messages = messages[-n*2-1:-1] if len(messages) >= n*2+1 else messages[:]

        user_msg = messages[-1].content if hasattr(messages[-1], "content") else str(messages[-1])
        history_list = []
        for m in messages[:-1]:
            role = "用户" if m.type == "human" or m.__class__.__name__ == "HumanMessage" else "助手"
            history_list.append(f"{role}: {m.content}")
        history_str = "\n".join(history_list)

        data = None
        for i in range(self.retry):
            raw = self.chain.invoke({"history": history_str, "user_input": user_msg})

            if raw.strip().startswith("```"):
                json_str = re.sub(r"^```(?:json)?", "", raw.strip())
                json_str = json_str.rstrip("`").strip()
            else:
                json_str = raw.strip()
                
            try:
                json_str = re.sub(r",\s*([\]}])", r"\1", json_str)
                data = json.loads(raw)
                break
            except json.JSONDecodeError:
                print(f"第 {i} 次信息抽取返回非 JSON")
                print(raw)
        else:
            print(f"信息抽取失败：未返回合法 JSON。")
            return
        
        with open("./log.txt", "a", encoding="utf-8") as f:
            f.write(f"add info: {data}\n")
        return data
